# -*- coding: utf-8 -*-
"""
@Time    : 2025/2/13 19:03 
@Author  : ZhangShenao 
@File    : 1.使用LangChain接入HuggingFace.py 
@Desc    : 使用LangChain接入HuggingFace
"""
import dotenv
import torch
from langchain_huggingface import HuggingFacePipeline, ChatHuggingFace

# 加载环境变量
dotenv.load_dotenv()

# 设置设备模式: cuda|cpu
device = 0 if torch.cuda.is_available() else -1

# 使用HuggingFace的Pipeline,创建一个LLM组件
llm = HuggingFacePipeline.from_model_id(
    model_id="Qwen/Qwen2.5-0.5B-Instruct",  # 指定HuggingFace上的模型ID
    task="text-generation",  # 指定任务类型
    device=device,  # 指定设备类型
    # 设置Pipeline参数
    pipeline_kwargs=dict(
        max_new_tokens=512,
        return_full_text=False,
    ),
)

# 调用ChatModel,获取结果
chat_model = ChatHuggingFace(llm=llm)
result = chat_model.invoke("写一首赞美秋天的五言绝句。")
print(result.content)
